#include "vector/tensor.h"

#include <stdint.h>
#include <numeric>
#include <execution>  // for std::execution::par
#include <stdexcept>  // for std::invalid_argument
#include <cmath> // For std::fabs

std::vector<std::vector<double>> Tensor::dot(const std::vector<std::vector<double>> &a, const std::vector<std::vector<double>> &b)
{
    auto n = a.size();
    auto m = b[0].size();
    std::vector<std::vector<double>> output(n, std::vector<double>(m));
    for(int64_t i = 0; i < n; i++)
    {
        for(int64_t j = 0; j < m; j++)
        {
            output[i][j] = 0;
            for(int64_t p = 0; p < b.size(); p++)
            {
                output[i][j] += a[i][p] * b[p][j];
            }
        }
    }
    return output;
}

std::vector<double> Tensor::dot(const std::vector<std::vector<double>> &a, const std::vector<double> &b)
{
    auto n = a.size();
    std::vector<double> output(n);
    for(int64_t i = 0; i < n; i++)
    {
        output[i] = 0;
        for(int64_t p = 0; p < b.size(); p++)
        {
            output[i] += a[i][p] * b[p];
        }
    }
    return output;
}
std::vector<std::vector<double>> Tensor::add(const std::vector<std::vector<double>> &a, const std::vector<std::vector<double>> &b)
{
    auto n = a.size();
    auto m = a[0].size();
    std::vector<std::vector<double>> output(n, std::vector<double>(m));
    for(int64_t i = 0; i < n; i++)
    {
        for(int64_t j = 0; j < m; j++)
        {
            output[i][j] = a[i][j] + b[i][j];
        }
    }
    return output;
}

std::vector<std::vector<double>> Tensor::transpose(const std::vector<std::vector<double>> &a)
{
    auto n = a.size();
    auto m = a[0].size();
    std::vector<std::vector<double>> output(a[0].size(), std::vector<double>(a.size()));
    for(int64_t i = 0; i < n; i++)
    {
        for(int64_t j = 0; j < m; j++)
        {
            output[j][i] = a[i][j];
        }
    }
    return output;
}

double Tensor::sum(const std::vector<double>& a) noexcept 
{
    return std::accumulate(a.begin(), a.end(), 0.0);
}

std::vector<double> Tensor::sum(const std::vector<std::vector<double>>& a, const int64_t dim) {
    if (a.empty()) {
        return {};
    }

    const size_t rows = a.size();
    const size_t cols = a[0].size();

    // 检查所有行是否具有相同的列数
    for (const auto& row : a) {
        if (row.size() != cols) {
            throw std::invalid_argument("All rows must have the same number of columns.");
        }
    }

    std::vector<double> result;

    if (dim == 0) {
        // 按列求和（并行优化）
        result.resize(cols, 0.0);
        for (size_t i = 0; i < rows; ++i) {
            // 使用 SIMD 或手动并行化（编译器可能自动优化）
            for (size_t j = 0; j < cols; ++j) {
                result[j] += a[i][j];
            }
        }
    } else if (dim == 1) {
        // 按行求和（使用 std::reduce + 并行策略）
        result.reserve(rows);
        for (const auto& row : a) {
            result.push_back(std::reduce(row.begin(), row.end(), 0.0));
        }
    } else {
        throw std::invalid_argument("Invalid dim value. Only 0 (column-wise) or 1 (row-wise) are supported.");
    }

    return result;
}

std::vector<std::vector<double>> Tensor::reshape(const std::vector<double>& input, int64_t n, int64_t m) {
    // 检查输入是否为空
    if (input.empty()) {
        throw std::invalid_argument("Input vector cannot be empty.");
    }

    // 检查形状参数是否合法
    if (n <= 0 || m <= 0) {
        throw std::invalid_argument("Shape dimensions must be positive.");
    }

    // 检查元素总数是否匹配
    if (static_cast<int64_t>(input.size()) != n * m) {
        throw std::invalid_argument("Input size does not match target shape.");
    }

    std::vector<std::vector<double>> output(n, std::vector<double>(m));

    // 按行优先顺序填充数据
    for (int64_t i = 0; i < n; ++i) {
        for (int64_t j = 0; j < m; ++j) {
            output[i][j] = input[i * m + j];
        }
    }

    return output;
}

std::vector<std::vector<double>> Tensor::multiply(const std::vector<std::vector<double>> &a, const std::vector<std::vector<double>> &b)
{
    std::vector<std::vector<double>> output = std::vector<std::vector<double>>(a.size(), std::vector<double>(a[0].size()));
    for(auto i = 0;i<a.size();i++)
    {
        for(auto j = 0;j<a[0].size();j++)
        {
            output[i][j] = a[i][j] * b[i][j];
        }
    }
    return output;
}

std::vector<std::vector<double>> Tensor::multiply(const std::vector<std::vector<double>> &a, double b)
{

    std::vector<std::vector<double>> output = std::vector<std::vector<double>>(a.size(), std::vector<double>(a[0].size()));
    for(auto i = 0;i<a.size();i++)
    {
        for(auto j = 0;j<a[0].size();j++)
        {
            output[i][j] = a[i][j] * b;
        }
    }
    return output;
}

std::vector<double> Tensor::multiply(const std::vector<double> &a, double b)
{
    std::vector<double> output = std::vector<double>(a.size());
    for(auto i = 0;i<a.size();i++)
    {
        output[i] = a[i] * b;
    }
    return output;
}

std::vector<std::vector<double>> Tensor::minus(const std::vector<std::vector<double>> &a, const std::vector<std::vector<double>> &b)
{
    std::vector<std::vector<double>> output = std::vector<std::vector<double>>(a.size(), std::vector<double>(a[0].size()));
    for(auto i = 0;i<a.size();i++)
    {
        for(auto j = 0;j<a[0].size();j++)
        {
            output[i][j] = a[i][j] - b[i][j];
        }
    }
    return output;
}

std::vector<double> Tensor::minus(const std::vector<double> &a, const std::vector<double> &b)
{
    std::vector<double> output = std::vector<double>(a.size());
    for(auto i = 0;i<a.size();i++)
    {
        output[i] = a[i] - b[i];
    }
    return output;
}


std::vector<std::vector<double>> Tensor::ones(int64_t row, int64_t col) {
    // 创建一个 row x col 的二维向量，所有元素初始化为 1.0
    return std::vector<std::vector<double>>(
        row, 
        std::vector<double>(col, 1.0)
    );
}

bool Tensor::equal(double a, double b, double epsilon)
{
    return std::abs(a - b) < epsilon;
}

bool Tensor::equal(const std::vector<std::vector<double>>& a, 
                    const std::vector<std::vector<double>>& b, 
                    double epsilon)
{
    if (a.size() != b.size()) return false;
    
    return std::equal(a.begin(), a.end(), b.begin(), 
        [epsilon](const auto& rowA, const auto& rowB) {
            if (rowA.size() != rowB.size()) return false;
            return std::equal(rowA.begin(), rowA.end(), rowB.begin(),
                [epsilon](double x, double y) { 
                    return equal(x, y, epsilon); 
                });
        });
}

bool Tensor::equal(const std::vector<double>& a, 
                  const std::vector<double>& b, 
                  double epsilon)
{
    // 首先检查大小是否相同
    if (a.size() != b.size()) {
        return false;
    }
    
    // 逐个比较元素
    for (size_t i = 0; i < a.size(); ++i) {
        if (std::fabs(a[i] - b[i]) > epsilon) {
            return false;
        }
    }
    
    return true;
}